import sys, os, time
from struct import pack, unpack
import struct
import socket, zlib, select
import datetime, math
from PIL import Image
from PIL import ImageDraw


#encoding-types
ENCODINGS=dict()
ENCODINGS['RAW']     = 0
ENCODINGS['COPY']    = 1
ENCODINGS['RRE']     = 2
ENCODINGS['CORRE']   = 4
ENCODINGS['HEXTILE'] = 5
ENCODINGS['ZLIB']    = 6
ENCODINGS['TIGHT']   = 7
ENCODINGS['ZLIBHEX'] = 8
ENCODINGS['ZRLE']    = 16
ENCODINGS['CURSOR']    = -239

#Special Keys
SpecialKeys=dict()
SpecialKeys['bs']     = 0xff08
SpecialKeys['tab']    = 0xff09
SpecialKeys['ret']    = 0xff0d
SpecialKeys['esc']    = 0xff1b
SpecialKeys['ins']    = 0xff63
SpecialKeys['del']    = 0xffff
SpecialKeys['home']   = 0xff50
SpecialKeys['end']    = 0xff57
SpecialKeys['pup']    = 0xff55
SpecialKeys['pdn']    = 0xff56
SpecialKeys['left']   = 0xff51
SpecialKeys['up']     = 0xff52
SpecialKeys['right']  = 0xff53
SpecialKeys['down']   = 0xff54
SpecialKeys['F1']     = 0xffbe
SpecialKeys['F2']     = 0xffbf
SpecialKeys['F3']     = 0xffc0
SpecialKeys['F4']     = 0xffc1
SpecialKeys['F5']     = 0xffc2
SpecialKeys['F6']     = 0xffc3
SpecialKeys['F7']     = 0xffc4
SpecialKeys['F8']     = 0xffc5
SpecialKeys['F9']     = 0xffc6
SpecialKeys['F10']    = 0xffc7
SpecialKeys['F11']    = 0xffc8
SpecialKeys['F12']    = 0xffc9
SpecialKeys['F13']    = 0xFFCA
SpecialKeys['F14']    = 0xFFCB
SpecialKeys['F15']    = 0xFFCC
SpecialKeys['F16']    = 0xFFCD
SpecialKeys['F17']    = 0xFFCE
SpecialKeys['F18']    = 0xFFCF
SpecialKeys['F19']    = 0xFFD0
SpecialKeys['F20']    = 0xFFD1
SpecialKeys['lshift'] = 0xffe1
SpecialKeys['shift']  = 0xffe1
SpecialKeys['rshift'] = 0xffe2
SpecialKeys['lctrl']  = 0xffe3
SpecialKeys['ctrl']   = 0xffe3
SpecialKeys['rctrl']  = 0xffe4
SpecialKeys['lmeta']  = 0xffe7
SpecialKeys['rmeta']  = 0xffe8
SpecialKeys['lalt']   = 0xffe9
SpecialKeys['alt']    = 0xffe9
SpecialKeys['ralt']   = 0xffea


SpecialKeys['SL']     = 0xFF14 #Scroll Lock
SpecialKeys['SR']     = 0xFF15 #Sys Request
SpecialKeys['NL']     = 0xFF7F #Num Lock
SpecialKeys['CL']     = 0xFFE5 #Caps Lock
SpecialKeys['Pause']  = 0xFF13
SpecialKeys['SupL']   = 0xFFEB #Super L
SpecialKeys['SupR']   = 0xFFEC #Super R
SpecialKeys['HypL']   = 0xFFED #Hyper L
SpecialKeys['HypR']   = 0xFFEE #Hyper R

SpecialKeys['KP0']    = 0xFFB0
SpecialKeys['KP1']    = 0xFFB1
SpecialKeys['KP2']    = 0xFFB2
SpecialKeys['KP3']    = 0xFFB3
SpecialKeys['KP4']    = 0xFFB4
SpecialKeys['KP5']    = 0xFFB5
SpecialKeys['KP6']    = 0xFFB6
SpecialKeys['KP7']    = 0xFFB7
SpecialKeys['KP8']    = 0xFFB8
SpecialKeys['KP9']    = 0xFFB9
SpecialKeys['KPEnter']= 0xFF8D

class RFBClient():
    def __init__(self, host, port, shared=1, password=False):
        self.ImageFormat='RGB'
        self.ImageDecode='RGBX'
        self.enc= list()
        self.color=True
        self.zlib=list()
        self.zlib.append(zlib.decompressobj())
        self.zlib.append(zlib.decompressobj())
        self.zlib.append(zlib.decompressobj())
        self.zlib.append(zlib.decompressobj())
        self.clipboard=""
        self.image = False
        self.cursor = False
        self.Complete = False
        self.Updates = 0
        self.shared=shared
        self.password=password
        self._buffer=bytearray()
        self.log=dict()
        self.log['levels']=["ERROR","INFO","DEBUG","DATA"]
        self.log['level']=0
        self.sock=dict()
        try:
            self.sock['fp'] = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock['fp'].connect((host, port))
            self.sock['fp'].setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            self.sock['fp'].setsockopt(socket.IPPROTO_TCP, socket.SO_RCVBUF, 65565*2)
            self.sock['fp'].setblocking(0)
            self.sock['epoll']=select.epoll()
            self.sock['epoll'].register(self.sock['fp'].fileno(), select.EPOLLIN)
        except:
            self.writeLog("Error, could not connect", level=0)
            sys.exit(1)
        self.read_data()
        self.initialize()

    def writeLog(self, data, level=2):
        if self.log['level'] >= level:
            print "-----------------\n%s: %s\n-----------------"%(self.log['levels'][level],data)

    def read_data(self, timeout=0.0):
        events = self.sock['epoll'].poll(timeout=timeout)
        for fn, event in events:
            if event == select.EPOLLIN and fn == self.sock['fp'].fileno():
                self._buffer.extend(self.sock['fp'].recv(4096))

    def write_data(self, data):
        self.sock['fp'].sendall(data)
        self.writeLog("Wrote data, %d"%len(data), level=3)

    def read_len(self, length):
        while len(self._buffer) < length:
            self.read_data(timeout=1.0)
        b = self._buffer[:length]
        self._buffer = self._buffer[length:]
        return str(b)

    def initialize(self):
        b = self.read_len(12)
        if b[:3] == 'RFB':
            b = b[3:]
            v = (b.strip()).split(".")
            if (int(v[0]), int(v[1])) not in [(3,3), (3,7), (3,8)]:
                print "Bad Protocol Version"
                sys.exit(1)
            b = b[12:]
            self.write_data('RFB 003.003\n')
            self.authclient()
        else:
            print "Does not seem to be an RFB port"
            sys.exit(1)

    def authclient(self):
        b = self.read_len(4)
        (auth,) = unpack("!I", b)
        if auth == 0:
            self.writeLog("Can not connect, bad auth", level=0)
            sys.exit(1)
        elif auth == 1:
            self.writeLog("No password needed, Continuing", level=1)
            self.write_data(pack("!B", self.shared))
        elif auth == 2 and self.password != False:
            self.writeLog("Can not do password auth currently", level=0)
            sys.exit(1)
        else:
            print "Bad Auth number %d"%auth
            sys.exit(1)
        b = self.read_len(24)
        (self.width, self.height, pixformat, namelen) = unpack("!HH16sI", b)
        (self.bpp, self.depth, self.bigendian, self.truecolor, self.redmax, self.greenmax, self.bluemax, self.redshift, self.greenshift, self.blueshift) = \
           unpack("!BBBBHHHBBBxxx", pixformat)
        self.bypp = self.bpp / 8        #calc bytes per pixel
        b = self.read_len(namelen)
        self.name = b
#        self.VNCstart()

    def update(self, timeout=0.1):
        if len(self._buffer) == 0:
            self.read_data(timeout=timeout)
        self.parsedata()

    def parsedata(self):
        while len(self._buffer) >= 1:
            b = self.read_len(1)
            (msgid,) = unpack("!B", b)
            if msgid == 0:
                self.image_update()
            elif msgid == 2:
                print "Got Bell"
            elif msgid == 3:
                self.ServerCutText()
            else:
                print "Getting bad data", msgid
                sys.exit(1)

    def image_update(self):
        b = self.read_len(3)
        (r,) = unpack("!xH", b)
        for i in range(r):
            self.Complete = False
            self.parse_rectangle()
        self.Complete = True
        self.Updates += 1

    def parse_rectangle(self):
        b = self.read_len(12)
        (x, y, width, height, encoding) = unpack("!HHHHi", b)
#        print "enc", encoding, x, y, width, height
        if encoding == ENCODINGS['COPY']:
            self.decodeCOPY(x, y, width, height)
        elif encoding == ENCODINGS['RAW']:
            self.decodeRAW( x, y, width, height)
        elif encoding == ENCODINGS['HEXTILE']:
            self.decodeHextile(x, y, width, height)
        elif encoding == ENCODINGS['CORRE']:
            print "CORRE"
            self.decodeRRE( x, y, width, height, CORRE=True)
        elif encoding == ENCODINGS['RRE']:
            self.decodeRRE(x, y, width, height)
        elif encoding == ENCODINGS['TIGHT']:
            self.decodeTight(x, y, width, height)
        elif encoding == ENCODINGS['ZLIB']:
            self.decodeRAW(x, y, width, height, compress = True)
        elif encoding == ENCODINGS['CURSOR']:
            self.decodePseudoCursor(x, y, width, height)
        else:
            print "ERROR: Bad Encoding", encoding
            sys.exit(1)

    def decodePseudoCursor(self, x, y, width, height):
        c = self.read_len(width*height*self.bypp)
        c = bytearray(c)
        bml = int(math.floor((width+7)/8.0)*height)
        bmw = int(math.floor((width+7)/8.0)*8)
        bm = self.read_len(bml)
        self.cursorBM=bytearray()
        for i in range(bml*8):
            l = 7-i%8
            cb = int(math.floor(i/8))
            bb = unpack("!B", bm[cb:cb+1])[0]
            bbb = ((bb >> l) & 0x01)
            if i % (bmw) < width:
                over = 0
                if i > 0:
                    over = int((bmw - width)*math.floor((i/bmw)))
                if bbb == 0:
                    self.cursorBM.append(0)
                else:
                    self.cursorBM.append(1)
        im = Image.fromstring(self.ImageFormat, (width,height), str(c), 'raw', self.ImageDecode)
        self.cursor = im
        self.cursorPOS=[x,y]
        
    # default protocol, pretty much required even if not set in encoding list
    def decodeRAW(self, x, y, width, height, compress=False):
        b = self.read_len(width*height*self.bypp)
        if compress == True:
            print "Got Compressed"
            b = zlib.decompress(b)
        if self.image == False:
            self.image = Image.new(self.ImageFormat, (self.width, self.height))
        try:
            tmp = Image.fromstring(self.ImageFormat, (width, height), b, 'raw', self.ImageDecode)
            self.image.paste(tmp, (x, y))
        except:
            print "Error",  x, y, width, height
            sys.exit(1)

    # can be very effective if used
    def decodeCOPY(self, x, y, width, height):
        b = self.read_len(4)
        if self.image == False:
            print "Self.image can not be false for CopyRect"
            sys.exit(1)
        (sx, sy) = unpack("!HH", b)
        cp = self.image.crop((sx,sy,sx+width,sy+height))
        cp.load()
        self.image.paste(cp, (x, y))

    # RRE and CORRE Encoding, Pretty common
    def decodeRRE(self, x, y, width, height, CORRE=False):
        b = self.read_len(4+self.bypp)
        (subrects,) = unpack("!I", b[:4])
        b = b[4:]
        if self.bypp == 4:
            (r,g,b,a) = unpack("!BBBB",b)
            c = (r,g,b,255)
        else:
            (c,) = unpack("!B",b[:1])
        im = Image.new(self.ImageFormat, (width, height), c)
        if subrects:
            if CORRE == False:
                b = self.read_len(((8+self.bypp)*subrects))
                subformat = "!HHHH"
            else:
                b = self.read_len(((4+self.bypp)*subrects))
                subformat = "!BBBB"
            b = bytearray(b) # Much faster to loop bytearray then str
            while len(b) > 0:
                if self.bypp == 4:
                    (sr,sg,sb,sa) = unpack("!BBBB", str(b[:4]))
                    b=b[4:]
                    c = (sr,sg,sb)
                else:
                    (sa,) = unpack("!B", str(b[:1]))
                    b=b[1:]
                    c = sa
                (sx, sy, swidth, sheight) = unpack(subformat, str(b[:struct.calcsize(subformat)]))
                b = b[struct.calcsize(subformat):]
                si = Image.new(self.ImageFormat, (swidth, sheight), c)
                im.paste(si, (sx,sy))
        if self.image == False:
            self.image = Image.new(self.ImageFormat, (self.width, self.height))
        self.image.paste(im, (x, y))


    # First Attempt at hexTile, it works but seems to take way to much CPU
    # probably because I am pasteing with PIL and should do everything
    # by hand till I make the final Image
    def decodeHextile(self, x, y, width, height):
        im = Image.new(self.ImageFormat, (width, height))
        tilesw = int(math.ceil((width)/16.0))
        tilesh = int(math.ceil((height)/16.0))
        if self.bypp == 4:
            bg=(0,0,0,0)
            fg=(0,0,0,0)
        else:
            bg=0
            fg=0
        for curY in range(tilesh):
            for curX in range(tilesw):
                setx = 16
                sety = 16
                subx = setx*curX
                suby = sety*curY
                b = self.read_len(1)
                subMask = unpack("!B", b)[0]
                if subx+16 > width:
                    setx=width-subx
                if suby+16 > height:
                    sety=height-suby
                tim = Image.new(self.ImageFormat, (setx,sety), bg)
                if subMask & 1: #Raw
                    b = self.read_len(setx*sety*self.bypp)
                    tim = Image.fromstring(self.ImageFormat, (setx,sety), b, 'raw', self.ImageDecode)
                else:
                    if subMask & 2:     #Background
                        b = self.read_len(self.bypp)
                        if self.bypp == 4:
                            (r,g,b,a) = unpack("!BBBB",b)
                            bg = (r,g,b,255)
                        else:
                            (c,) = unpack("!B",b)
                            bg = c
                    if subMask & 4:     #Foreground
                        b = self.read_len(self.bypp)
                        if self.bypp == 4:
                            fg = unpack("!BBBB",b)
                        else:
                            (c,) = unpack("!B",b)
                            fg = c
                    tim = Image.new(self.ImageFormat, (setx,sety), bg)
                    if subMask & 8:     #Subrects
                        b = self.read_len(1)
                        sr = unpack("!B", b)[0]
                        for subrect in range(sr):
                            fill=fg
                            if subMask & 16:
                                b = self.read_len(self.bypp)
                                if self.bypp == 4:
                                    (r,g,b,a) = unpack("!BBBB",b)
                                    fill=(r,g,b,255)
                                else:
                                    (c,) = unpack("!B",b)
                                    fill = c
                            b = self.read_len(2)
                            (subPos, subD) = unpack("!BB", b)
                            rsy = subPos & 0x0F
                            rsx = subPos >> 4
                            rsh = subD & 0x0F
                            rsw = subD >> 4
                            stim = Image.new(self.ImageFormat, (rsw+1,rsh+1), fill)
                            tim.paste(stim, (rsx,rsy))
                im.paste(tim, (curX*16, curY*16))
        if self.image == False:
            self.image = Image.new(self.ImageFormat, (self.width, self.height))
        self.image.paste(im, (x, y))


    # Tight encoding, does not work still figuring it out
    def decodeTight(self, x, y, width, height):
        print "Tight", x, y, width, height
        b = self.read_len(1)
        (c,) = unpack("!B",b)

        if (c >> 4) == 8:
            print "Fill"
            if self.bypp == 4:
                self.read_len(3)
            else:
                self.read_len(self.bypp)
        elif (c >> 4) == 9:
            print "Jpeg"
        elif (c >> 4) < 8:
            print "Basic"
            if c & 64:
                (f,) = unpack("!B", self.read_len(1))
                if f == 1:
                    print "Palette filter", f
#                    (pf,) = unpack("!B", self.read_len(1))
#                    print pf
                    x = y = z = 0
                    x = unpack("!B", self.read_len(1))[0]
                    if x > 127:
                        y = unpack("!B", self.read_len(1))[0]
                    if y > 127:
                        z = unpack("!B", self.read_len(1))[0]
                    l = (z<< 14) | ((y & 0x7f) << 7) | (x & 0x7f)
                    print l
                    self.read_len(l)
                    
        print bin(c)
        

    # ---  ZRLE Encoding
    def decodeZRLE(self, block):
        pass

    def ServerCutText(self):
        b = self.read_len(7)
        (length, ) = unpack("!xxxI", b)
        b = self.read_len(length)
        self.clipboard = b
        
    def setPixelFormat(self, color=True):
        if color == True:
            bpp=32
            depth=24
            self.ImageFormat = 'RGB'
            self.ImageDecode = 'RGBX'
        else:
            bpp=8
            depth=1
            self.ImageFormat = 'L'
            self.ImageDecode = "L"
        bigendian=0
        truecolor=1
        redmax=255
        greenmax=255
        bluemax=255
        redshift=0
        greenshift=8
        blueshift=16
        self.writeLog("Setting Pix Format")
        pixformat = pack("!BBBBHHHBBBxxx", bpp, depth, bigendian, truecolor, redmax, greenmax, bluemax, redshift, greenshift, blueshift)
        self.write_data(pack("!Bxxx16s", 0, pixformat))
        self.bpp, self.depth, self.bigendian, self.truecolor = bpp, depth, bigendian, truecolor
        self.redmax, self.greenmax, self.bluemax = redmax, greenmax, bluemax
        self.redshift, self.greenshift, self.blueshift = redshift, greenshift, blueshift
        self.bypp = self.bpp / 8        

    def setEncodings(self, l):
        l.append(ENCODINGS['CURSOR'])
        self.writeLog("Setting Encodings")
        self.write_data(pack("!BxH", 2, len(l)))
        for encoding in l:
            self.write_data(pack("!i", encoding))

    def setContinuous(self, enable=1):
        self.writeLog("Set Continuous")
        self.write_data(pack("!BBHHHH", 150,enable,0,0,self.width,self.height))

    def imageUpdateRequest(self, inc=0):
        self.writeLog("Request Update")
        self.write_data(pack("!BBHHHH", 3, inc, 0, 0, self.width, self.height))

    def sendString(self, s):
        for i in s:
            time.sleep(.01)
            self.write_data(pack("!BBxxI",4,1, ord(i)))
            time.sleep(.01)
            self.write_data(pack("!BBxxI",4,0, ord(i)))
        time.sleep(.01)


    def pointerEvent(self, x, y, buttonmask=0):
        if x > self.width or y > self.height:
            print "mouse error"
            return False
        self.write_data(pack("!BBHH", 5, buttonmask, x, y))

    def clientCutText(self, message):
        self.write_data(pack("!BxxxI", 6, len(message)) + message)
    def keyEvent(self, key, down=0):
        self.write_data(pack("!BBxxI",4,down, key))

    def sendSpecial(self, c):
        l = c.split("+")
        key = None
        p = None
        for i in l:
            if i in SpecialKeys:
                key = SpecialKeys[i]
            elif i == "1" or i == "0":
                p = int(i)
            if p != None and key != None:
                self.keyEvent(key, down=p)
                key = None


    def VNCstart(self):
        if len(self.enc) == 0:
            self.setEncodings([ENCODINGS['COPY'], ENCODINGS['RRE']])
        self.setPixelFormat(color=self.color)
        self.imageUpdateRequest()
        self.update()
        



